{
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.6-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "source": [
    "给定一个整数数组 nums，返回区间和在 `[lower, upper]` 之间的个数，包含 lower 和 upper。\n",
    "区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。\n",
    "\n",
    "说明:\n",
    "最直观的算法复杂度是 O(n2) ，请在此基础上优化你的算法。\n",
    "```\n",
    "示例:\n",
    "\n",
    "输入: nums = [-2,5,-1], lower = -2, upper = 2,\n",
    "输出: 3 \n",
    "解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2。\n",
    "```\n",
    "来源：力扣（LeetCode）\n",
    "链接：https://leetcode-cn.com/problems/count-of-range-sum\n",
    "著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "除了暴力，我根本想不到其他的算法，"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "source": [
    "暴力的复杂度为$O(n)$"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_range_sum(nums, lower, upper):\n",
    "    def merge_count(arr):\n",
    "        n = len(arr)\n",
    "        if n <= 1:\n",
    "            return 0\n",
    "        mid = n//2\n",
    "        n1 = arr[:mid]\n",
    "        n2 = arr[mid:]\n",
    "        print(n1, n2)\n",
    "        cnt = merge_count(n1) + merge_count(n2)\n",
    "        print(\"cnt=\", cnt)\n",
    "        l, r = 0, 0\n",
    "        for v in n1:\n",
    "            while l < len(n2) and n2[l]-v < lower:\n",
    "                l+=1\n",
    "            while r < len(n2) and n2[r]-v<= upper:\n",
    "                r+=1\n",
    "            cnt += r -l\n",
    "        \n",
    "        p1, p2 = 0, 0\n",
    "        for i in range(len(arr)):\n",
    "            if p1 < len(n1) and (p2 == len(n2) or n1[p1] <= n2[p2]):\n",
    "                arr[i] = n1[p1]\n",
    "                p1 += 1\n",
    "            else:\n",
    "                p2 += 1\n",
    "        \n",
    "        return cnt\n",
    "    \n",
    "    prefix_sum = [0]\n",
    "    for i in range(len(nums)):\n",
    "        prefix_sum.append(prefix_sum[i] + nums[i])\n",
    "    print(prefix_sum)\n",
    "    return merge_count(prefix_sum)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[0, -2, 3, 2]\n[0, -2] [3, 2]\n[0] [-2]\ncnt= 0\n[3] [2]\ncnt= 0\ncnt= 2\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "metadata": {},
     "execution_count": 24
    }
   ],
   "source": [
    "count_range_sum([-2, 5, -1], -2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}